from fastapi import UploadFile
from starlette.responses import FileResponse
from sqlalchemy import and_
from datetime import datetime

import os
from app.dao import Session
from app.config import DEFAULT_KNOWLEDGE_HUB_PATH
from app.models import Group,User, Hub, GroupDetail,File
from app.parameter import  QueryFilesParam,QueryHubsParam,CreateHubParam,DeleteHubParam,QueryHubParam
from app.utils import Response, signatureDecode,signatureEncode
from app.constant import ERRORCODE
from store import VectorStore
from chunks import UnstructuredChunk
from embeddings import embeddingPool
from utils import DEBUG,ERROR

from threading import Thread

class FileChunkTask(Thread):
    
    def __init__(self, path: str,filename: str, hub_id: int, file_id: id, chunkSize: int, overlap: int, embed):
        super().__init__()
        self.path = path
        self.filename = filename
        self.hub_id = hub_id
        self.file_id = file_id
        self.chunkSize = chunkSize
        self.overlap = overlap
        self.embed = embed
        
    def processFile(self):
        try:
            chunk = UnstructuredChunk(self.chunkSize, self.overlap)
            chunk.load(self.path)
            res = chunk.chunk()
            embeddings = [self.embed(i) for i in res[0]]
            store = VectorStore(self.hub_id)
            store.load()
            store.insert(self.file_id, self.filename, res[1], res[0], embeddings)
            
        # 文档切分入库失败，修改状态为 2
        except Exception as e:
            ERROR(e)
            with Session() as session:
                session.query(File).filter(File.id == self.file_id).update({"state": 2})
                session.commit()    
            return 
        
        # 文档切分入库失败，修改状态为 1
        with Session() as session:
            session.query(File).filter(File.id == self.file_id).update({"state": 1})
            session.commit()
        
    def run(self):
        self.processFile()



class FileService():
    @staticmethod
    def judgeGroup(userId :int, groupId :int, session):
        return session.query(Group).filter(and_(Group.user_id == userId, Group.group_id == groupId)).first()

    @staticmethod
    def judgeHub(userId :int, hubId :int, session):
        return session.query(Hub, Group).join(Group, Hub.group_id == Group.group_id).filter(and_(Hub.id == hubId, Group.user_id == userId)).first()
    
    @staticmethod
    def queryFiles(queryFilesParam :QueryFilesParam, userId: int):
        try:
            with Session() as session:
                group = FileService.judgeHub(userId, queryFilesParam.hubId, session)
                if group == None:
                    return Response(ERRORCODE.REQUEST_NOT_FOUND_ERROR, message="not found group!")
                
                files = session.query(File).filter(File.hub_id == queryFilesParam.hubId).all()
                files = [file.to_dict() for file in files]
        
                return Response(ERRORCODE.RESPONSE_SUCCESS, files)
        except Exception as e:
            ERROR(e)
            return Response(ERRORCODE.RESPONSE_DATABASE_OPERATION_ERROR, None)
    
    @staticmethod
    async def createFile(file :UploadFile, hubId: int ,userId: int, embed = None):
        filepath = ""
        fileId = 0
        chunkSize = 0
        overlap = 0
        
        try:
            with Session() as session:
                DEBUG(hubId)
                res = FileService.judgeHub(userId, hubId, session)
                if res == None:
                    return Response(ERRORCODE.REQUEST_NOT_FOUND_ERROR, message="not found group!")
            
                DEBUG(res)
                group = res[1].to_dict()
                DEBUG(group)
                if group["permission"] < 1:
                    return Response(ERRORCODE.REQUEST_HAVE_NO_PERMISSION)
                
                hub = res[0].to_dict()
                chunkSize = hub["chunk_size"]
                overlap = hub["overlap"]
                
                queryFile = session.query(File).filter(and_(File.hub_id == hubId, File.file_name == file.filename)).first()
                if queryFile != None:
                    return Response(ERRORCODE.REQUEST_ERROR,  message="The file {} already exists in hub!".format(file.filename))
                
                filepath = DEFAULT_KNOWLEDGE_HUB_PATH + "/" + str(group["group_id"]) +"/" \
                    + str(hubId) 
                    
                if not os.path.isdir(filepath):
                    os.makedirs(filepath)
                    
                filepath = filepath + "/"+ file.filename   
                DEBUG(filepath)
                contents = await file.read()   
                with open(filepath, "wb") as f:
                    f.write(contents)
                    
                fileEntry = File(hub_id = hubId, user_id = userId, file_name = file.filename)
                session.add(fileEntry)
                session.flush()
                fileId = fileEntry.id
                session.commit()

            if embed == None:
                p = FileChunkTask(filepath, file.filename, hubId, fileId, chunkSize, overlap, embeddingPool.get())
            else:
                p = FileChunkTask(filepath, file.filename, hubId, fileId, chunkSize, overlap, embed)
            p.start()
            return Response(ERRORCODE.RESPONSE_SUCCESS)
            
        except Exception as e:
            raise e
            ERROR(e)
            return Response(ERRORCODE.RESPONSE_DATABASE_OPERATION_ERROR)


    ## no finish
    @staticmethod
    def reloadFile(fileId :int, hubId: int, userId :int, embed = None):

        with Session() as session:
            result = session.query(File, Hub).join(Hub, File.hub_id == Hub.id).filter(and_(and_(File.id == fileId, File.user_id == userId), File.state == 2)).first()
            DEBUG(result)
            if result == None:
                return Response(ERRORCODE.REQUEST_NOT_FOUND_ERROR, "未找到文件或者没有权限")
            session.query(File).filter(File.id == fileId).update({"state": 0})
            file = result[0]
            file = file.to_dict()
            hub = result[1]
            hub = hub.to_dict()
            
            filename = file["file_name"]
            hub_id = hub["id"]
            group_id = hub["group_id"]
            chunkSize = hub["chunk_size"]
            overlap = hub["overlap"]
            
            filepath = DEFAULT_KNOWLEDGE_HUB_PATH + "/" + str(group_id) +"/" \
                    + str(hubId)
            if not os.path.isdir(filepath):
                    os.makedirs(filepath)
                    
            filepath = filepath + "/"+ filename
                
            if embed == None:
                p = FileChunkTask(filepath, filename, hub_id, fileId, chunkSize, overlap, BgeEmbedding())
            else:
                p = FileChunkTask(filepath, filename, hub_id, fileId, chunkSize, overlap, embed)
            
            p.start()
            session.commit()
            
            return Response(ERRORCODE.RESPONSE_SUCCESS)
    
    @staticmethod
    async def downloadFile(userId :int, hubId :int, fileId :int):
        try:
            with Session() as session:
                # 查询
                res = FileService.judgeHub(userId, hubId, session)
                
                if res == None:
                    return Response(ERRORCODE.REQUEST_HAVE_NO_PERMISSION)
                group = res[1].to_dict()
                file = session.query(File).filter(and_(File.id == fileId, File.hub_id == hubId)).first()

                if file == None:
                    return Response(ERRORCODE.REQUEST_NOT_FOUND_ERROR, message="Not find the file!")
                file = file.to_dict()
                fileName =  file["file_name"] 
                filepath = DEFAULT_KNOWLEDGE_HUB_PATH + "/" + str(group["group_id"]) + "/" + str(hubId) + "/" + fileName
                DEBUG(filepath)
                return FileResponse(path=filepath,filename=fileName)
        except Exception as e:
            ERROR(e)
            return Response(ERRORCODE.RESPONSE_DATABASE_OPERATION_ERROR, None)

    @staticmethod
    def deleteFile(deleteHubParam :DeleteHubParam, userId: int):
        try:
            with Session() as session:
                # 查询
                hub = session.query(Hub)\
                    .filter(and_(Hub.id == deleteHubParam.hubId, Hub.user_id == deleteHubParam.userId)).delete()
                if hub == None:
                    return Response(ERRORCODE.REQUEST_HAVE_NO_PERMISSION)
                print(hub)
                session.commit()
                return Response(ERRORCODE.RESPONSE_SUCCESS,hub)
        except Exception as e:
            ERROR(e)
            return Response(ERRORCODE.RESPONSE_DATABASE_OPERATION_ERROR, None)

    @staticmethod
    def queryFile(queryHubParam :QueryHubParam, userId):
        try:
            with Session() as session:
                res = HubService.judgeHub(userId, queryHubParam.hubId, session)
                DEBUG(res)
                if res == None:
                    return Response(ERRORCODE.REQUEST_HAVE_NO_PERMISSION)
                
                data = session.query(Hub).join(GroupDetail, Hub.group_id == GroupDetail.id).\
                    join(User, Hub.user_id == User.id).filter(Hub.id == queryHubParam.hubId)\
                        .with_entities(Hub.hub_name, Hub.description, Hub.create_at,User.nickname, User.avatar,\
                                       GroupDetail.group_name).first();
                DEBUG(data)
                
                res = {}
                res['hub_name'] = data[0]
                res['description'] = data[1]
                res['create_at'] = str(data[2])
                res['nickname'] = data[3]
                res['avatar'] = data[4]
                res['group_name'] = data[5]
                return Response(ERRORCODE.RESPONSE_SUCCESS, res)
        except Exception as e:
            ERROR(e)
            return Response(ERRORCODE.RESPONSE_DATABASE_OPERATION_ERROR, None)